library(tidyverse)
library(ggpubr)
library(LDlinkR)
library(LDheatmap)

merge_depth <- function(string){
  string |>
    str_split_1(',') |>
    as.numeric() |>
    sum()
}

g2r_position <- 106204113

# published filtered vcf -------------
vindi_vcf <- read_delim('Archive/covid19/data/14.106070000-106260000.chr14_mq25_mapab100.vcf.gz', comment = '##') |>
  select(-c(INFO, FILTER, ID))

vindi_vcf_tidy <- vindi_vcf |>
  separate(col = Vindija33.19, into = c('genotype', 'depth', 'A', 'C', 'G', 'T', 'phred_probability', 'genotype_quality'), sep = ':')

# snpAD-called unfiltered vcf -----------------
unfilter_vindi_vcf <- read_delim('Archive/covid19/data/snpAD_unfiltered.vcf.gz', comment = '##') |>
  separate(col = Vindija33.19, into = c('genotype', 'depth', 'A', 'C', 'G', 'T', 'phred_probability', 'genotype_quality'), sep = ':') |>
  select(-c(INFO, FILTER, ID, FORMAT, genotype_quality)) |>
  mutate(depth = as.numeric(depth))

unfilter_vindi_vcf |>
  ggplot(aes(QUAL)) +
  geom_boxplot()

median(unfilter_vindi_vcf$depth)

# plot sequencing depth -----------
unfilter_vindi_vcf |>
  ggplot(aes(depth)) +
  geom_boxplot() +
  labs_pubr() +
  theme(axis.text.y = element_blank()) +
  labs(x = 'sequencing depth')

g1 <- last_plot()

unfilter_vindi_vcf |>
  ggplot(aes(as.numeric(depth))) +
  geom_density() +
  geom_vline(xintercept = 12, linetype = 'dashed') +
  labs_pubr() +
  labs(x = '', y = 'frequency of reads')

g2 <- last_plot()

patchwork::plot_layout(g2 / g1)

# samtools depth raw as igv --------
samtools_depth <- read_delim('Archive/covid19/data/vindi_depth.txt', col_names = c('chrom','position','depth'))

depth_median <- median(samtools_depth$depth)
mean(samtools_depth$depth)

samtools_depth |>
  ggplot(aes(depth)) +
  geom_boxplot() +
  labs_pubr() +
  theme(axis.text.y = element_blank()) +
  labs(x = 'sequencing depth')

g1 <- last_plot()

samtools_depth |>
  ggplot(aes(as.numeric(depth))) +
  geom_density() +
  geom_vline(xintercept = depth_median, linetype = 'dashed') +
  labs_pubr() +
  labs(x = '', y = 'frequency of reads') +
  labs(title = str_glue('Median depth: {depth_median}'))

g2 <- last_plot()

patchwork::plot_layout(g2 / g1)

unfilter_vindi_vcf |> count(genotype)

CDX_other_fst <- read_csv('Archive/covid19/results/CDX_vs_non-EAS_Fst.csv')

high_fst_threshold <- quantile(CDX_other_fst$Fst, 0.99)

YangTY <- read_delim('Archive/covid19/data/YangTY_summary.txt')

fst_with_depth <- CDX_other_fst |>
  left_join(unfilter_vindi_vcf, join_by(position == POS)) |>
  filter(!is.na(depth)) |>
  rowwise() |>
  mutate(across(9:12, merge_depth)) |>
  ungroup()

fst_with_depth |>
  ggplot(aes(position, Fst)) +
  geom_point()

good_depth_fst <- YangTY |>
  left_join(CDX_other_fst) |>
  filter(Fst >= high_fst_threshold & Vindija_depth > 11)

# plot fst with good depth -------------
CDX_other_fst |>
  filter(position > 10607e4) |>
  ggplot(aes(position, Fst)) +
  geom_point(alpha = 0.1) +
  geom_point(data = good_depth_fst, color = 'red') +
  geom_point(data = filter(CDX_other_fst, position == g2r_position), color = 'orange', size = 5, shape = "circle open") +
  geom_hline(yintercept = high_fst_threshold, linetype = 'dashed') +
  ylab('Fst (CDX vs non-EAS)') +
  theme_pubr() +
  labs_pubr()

fst_with_depth |>
  filter(position %in% good_depth_fst$position)

# calculate LD D' & R2
api_token <- Sys.getenv("LDLINK_TOKEN")

find_ctrl_sites <- function(candidates, n){
  CDX_other_fst |>
    filter(position < min(candidates) | position > max(candidates)) |>
    slice_sample(n = n) |>
    pull(position)
}

set.seed(42)

snp_matrix <- good_depth_fst$position |>
  find_ctrl_sites(20) |>
  append(c(good_depth_fst$position, g2r_position)) |>
  sort() |>
  map(\(x)str_glue('chr14:{x}')) |>
  SNPclip(pop = 'EAS', token = api_token)

calc_ld_stat <- function(rs_list) {
  r2_stat <- rs_list |>
    LDmatrix(pop = 'EAS',
             token = api_token) |>
    filter(RS_number == 'rs117518546') |>
    pivot_longer(cols = where(is.numeric),
                 names_to = 'linked_snp',
                 values_to = 'r2') |>
    select(c(linked_snp, r2))
  
  rs_list |>
    LDmatrix(pop = 'EAS',
             token = api_token,
             r2d = "d") |>
    filter(RS_number == 'rs117518546') |>
    pivot_longer(cols = where(is.numeric),
                 names_to = 'linked_snp',
                 values_to = 'Dprime') |>
    select(c(linked_snp, Dprime)) |>
    left_join(r2_stat)
}

g2r_ld_stat <- calc_ld_stat(snp_matrix$RS_Number)

# LD plot (heatmap) ---------
plot_Dprime_map <- function(whole_list, block_pos, plot_title) {
  ld_matrix <- whole_list |>
    LDmatrix(pop = 'EAS',
             r2d = 'd',
             token = api_token) |>
    column_to_rownames('RS_number') |>
    mutate(across(everything(), \(x)replace_na(x, 0L)))
  
  pos_matrix <- ld_matrix |>
    colnames() |>
    SNPclip(pop = 'EAS', token = api_token) |>
    mutate(numeric_pos = str_remove(Position, 'chr14:') |>
             as.numeric())
  
  block_index1 <- pos_matrix$numeric_pos %in% block_pos |>
    detect_index(isTRUE)
  
  block_index2 <- pos_matrix$numeric_pos %in% block_pos |>
    detect_index(isTRUE, .dir = 'backward')
  
  ld_matrix |>
    as.matrix() |>
    LDheatmap(genetic.distances = pos_matrix$numeric_pos,
              flip = TRUE,
              LDmeasure = "D'",
              title = plot_title,
              SNP.name = 'rs117518546',
              color = colorRampPalette(c('red','yellow','white'))(20)) |>
    LDheatmap.highlight(block_index1, block_index2, lwd = 3)
}

snp_matrix$RS_Number |>
  plot_Dprime_map(block_pos = good_depth_fst$position,
                  plot_title = '9 high-quality SNP in IGHG1-G396R LD block')

max(good_depth_fst$position) - min(good_depth_fst$position)

good_depth_fst |>
  ggplot(aes(position, Vindija_depth)) +
  geom_col() +
  labs_pubr()

final_good_snp_stat <- fst_with_depth |>
  select(c(position, snpAD_depth = depth)) |>
  right_join(good_depth_fst) |>
  mutate(chrpos = str_glue('chr14:{position}')) |>
  left_join(snp_matrix[1:2], join_by(chrpos == Position)) |>
  left_join(g2r_ld_stat, join_by(RS_Number == linked_snp))


# plot subpop MAF -------------
maf_sub <- read_csv('Archive/covid19/results/1KGP_subpop_MAF.csv')

sub_color = c(colorRampPalette(c('red','orange'))(6),
              colorRampPalette(c('grey60','black'))(20))

maf_sub |>
  filter(POS %in% good_depth_fst$position) |>
  mutate(pop = fct_relevel(pop, 'CDX', 'KHV', 'CHS', 'CHB', 'BEB', 'JPT')) |>
  ggplot(aes(pop, MAF, fill = pop)) +
  geom_col() +
  facet_wrap(~POS) +
  scale_fill_manual(values = sub_color) +
  theme(axis.text.x = element_text(angle = 90, hjust = 1, vjust = .5)) +
  scale_x_discrete(breaks = NULL) +
  labs(fill = 'population')

maf_sub |>
  pivot_wider(names_from = 'pop', values_from = 'MAF', names_prefix = 'MAF_') |>
  right_join(final_good_snp_stat, join_by(POS == position)) |>
  relocate(POS, RS_Number, Vindija_depth, snpAD_depth, Dprime, r2) |>
  write_csv('Archive/covid19/results/linked_above11reads_snp.csv')
